package com.thermal.thermal_inspection.controller;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.thermal.thermal_inspection.entity.CaseInfo;
import com.thermal.thermal_inspection.service.CaseInfoService;
import io.reactivex.Flowable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@RestController
@RequestMapping("/ai")
public class AIController {

    private static final Logger logger = LoggerFactory.getLogger(AIController.class);
    private final Generation generation = new Generation();

    @Autowired
    private CaseInfoService caseInfoService;

    /**
     * 流式接口 - Server-Sent Events (SSE)
     */
    @GetMapping(value = "/chat/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamChat(@RequestParam String question, @RequestParam int caseId) {
        List<CaseInfo> caseInfos = caseInfoService.getByCaseId(caseId);
        String prompt = buildChatPrompt(caseInfos);
        String questionForAI = prompt + "\n请结合以上聊天内容，继续回答我下面的问题：\n用户: " + question;
        int caseInfoId;
        try {
            // ① 先保存用户提问（answerInfo设为null）
            CaseInfo caseInfo = new CaseInfo();
            caseInfo.setCaseId(caseId);
            caseInfo.setQuestionInfo(question);
            caseInfo.setAnswerInfo(null);
            caseInfoId = caseInfoService.addQuestion(caseInfo); // 返回自增主键ID
        } catch (Exception e) {
            logger.error("插入问题失败", e);
            return Flux.just(ServerSentEvent.<String>builder()
                    .event("error")
                    .data("问题入库失败：" + e.getMessage())
                    .build());
        }
        try {
            Message userMsg = Message.builder()
                    .role(Role.USER.getValue())
                    .content(questionForAI)
                    .build();

            GenerationParam param = buildGenerationParam(userMsg);
            Flowable<GenerationResult> flowable = generation.streamCall(param);

            // 用于收集完整回答
            List<String> collectedChunks = new ArrayList<>();

            return Flux.from(flowable)
                    .map(result -> {
                        String chunk = result.getOutput().getChoices().get(0).getMessage().getContent();
                        chunk = chunk.replace(" ", "&nbsp;"); // 仅用于HTML场景
                        collectedChunks.add(chunk); // 收集内容
                        return ServerSentEvent.<String>builder()
                                .data(chunk)
                                .build();
                    })
                    .doOnComplete(() -> {
                        // ② AI回答完成后，更新到数据库
                        String fullAnswer = String.join("", collectedChunks);
                        boolean updated = caseInfoService.updateAnswer(caseInfoId, fullAnswer);
                        if (!updated) {
                            logger.warn("更新回答失败，caseInfoId: {}", caseInfoId);
                        }
                    })
                    .onErrorResume(e -> {
                        logger.error("AI回答流式异常", e);
                        return Flux.just(ServerSentEvent.<String>builder()
                                .event("error")
                                .data("AI服务异常：" + e.getMessage())
                                .build());
                    });

        } catch (Exception e) {
            logger.error("流式初始化失败", e);
            return Flux.just(ServerSentEvent.<String>builder()
                    .event("error")
                    .data("初始化失败：" + e.getMessage())
                    .build());
        }
    }

    // 构建参数
    private GenerationParam buildGenerationParam(Message userMsg) {
        return GenerationParam.builder()
                .apiKey("sk-867e2d2f25a342a0a7164fcf15583ef4") // 建议读取配置文件
                .model("qwen-plus")
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }

    public String buildChatPrompt(List<CaseInfo> caseInfoList) {
        if (caseInfoList == null || caseInfoList.isEmpty()) {
            return "";
        }

        StringBuilder prompt = new StringBuilder();

        // 添加提示词
        prompt.append("以下是我与AI之前的聊天记录，请参考这些历史内容继续回答我的新问题：\n\n");

        // 添加第一条记录
        CaseInfo first = caseInfoList.get(0);
        prompt.append("用户: ").append(first.getQuestionInfo()).append("\n");
        prompt.append("AI: ").append(first.getAnswerInfo()).append("\n\n");

        // 添加最后三条记录，避免重复第一条
        int size = caseInfoList.size();
        int start = Math.max(1, size - 3); // 从第二条开始
        for (int i = start; i < size; i++) {
            CaseInfo info = caseInfoList.get(i);
            prompt.append("用户: ").append(info.getQuestionInfo()).append("\n");
            prompt.append("AI: ").append(info.getAnswerInfo()).append("\n\n");
        }

        return prompt.toString().trim();
    }
}
